'''
Created on 08.06.2012

@author: Julia
'''

from __future__ import division
import numpy as np
import matplotlib.pyplot as plt

sz=0.5*np.array([[1, 0], [0, -1]])
sp=np.array([[0, 1], [0, 0]])
sm=np.array([[0, 0], [1, 0]])
m=4
n0=3

class block():
    def __init__(self, n0):
        """
        H11: Hamiltonian of the block
        Szl, Spl, Sml: Operators for the rightest site of the block
        Sz0, Sp0, Sm0: Operators for the leftest site of the block
        """
        self.H11=np.zeros((2**n0,2**n0), dtype=float)
        for i in range(n0-1):
            self.H11+=np.kron(np.kron(np.eye(2**i),np.kron(sz,sz)+0.5*np.kron(sp,sm)+0.5*np.kron(sm,sp)),np.eye(2**(n0-i-2)))
        self.dim=2**n0
        self.szl=np.kron(np.eye(self.dim/2),sz)
        self.spl=np.kron(np.eye(self.dim/2),sp)
        self.sml=np.kron(np.eye(self.dim/2),sm)
        self.sz0=np.kron(sz,np.eye(self.dim/2))
        self.sp0=np.kron(sp,np.eye(self.dim/2))
        self.sm0=np.kron(sm,np.eye(self.dim/2))
        self.sites=n0
        
       
        
def T(b1, b2):
    """
    Interaction between two blocks b1 (right) and b2 (left)
    """
    T12=np.kron(b1.szl, b2.sz0)+0.5*np.kron(b1.spl, b2.sm0)+0.5*np.kron(b1.sml, b2.sp0)
    return T12

class system():
    def __init__(self):
        self.b1=block(n0)
        self.b2=block(1)
        self.b3=block(1)
        self.b4=block(n0)
    
    
    def get_gs(self):
        """
        builds up the total Hamiltonian for the 4 blocks b1, b2, b3 and b4
        gets its lowest eigenvalue as self.energy, lowest eigenvector as self.psi
        """
        H11_tot=np.kron(np.kron(np.kron(self.b1.H11, np.eye(self.b2.dim)), np.eye(self.b3.dim)), np.eye(self.b4.dim))+np.kron(np.kron(np.kron(np.eye(self.b1.dim), self.b2.H11), np.eye(self.b3.dim)), np.eye(self.b4.dim))+np.kron(np.kron(np.kron(np.eye(self.b1.dim), np.eye(self.b2.dim)), self.b3.H11), np.eye(self.b4.dim))+np.kron(np.kron(np.kron(np.eye(self.b1.dim), np.eye(self.b2.dim)), np.eye(self.b3.dim)), self.b4.H11)
        T_tot=np.kron(T(self.b1, self.b2), np.eye(self.b3.dim*self.b4.dim))+np.kron(np.kron(np.eye(self.b1.dim), T(self.b2, self.b3)), np.eye(self.b4.dim))+np.kron(np.eye(self.b1.dim*self.b2.dim), T(self.b3, self.b4))
        H_tot=H11_tot+T_tot
        #print "get_gs: H_tot(1, :)", H_tot[0, :]
        [u, v]=np.linalg.eig(H_tot)
        idx=u.argsort()
        u=u[idx]
        v=v[:, idx]
        #print "v davor:", v[:, 0]
        #print "abs:", v[:, 0].sum()
        if v[:, 0].sum()<0:
            v[:, 0]=v[:, 0]*(-1)
        self.psi=v[:, 0]
        self.energy=u[0]
        #print "psi:", self.psi
        #print 'u for psi:', u.T
    
    def basis_left(self, new):
        """
        adds a site to the left of block self, thereby creates the new block "new"
        builds the density matrix from self.psi, takes its m highest eigenvalues and transforms the
        block using the corresponding eigenvectors
        """
        dim_sys=self.b1.dim*self.b2.dim
        dim_env=self.b3.dim*self.b4.dim
        #print 'psi, dim_sys, dim_env', self.psi, dim_sys, dim_env
        psi_matrix=self.psi.reshape(dim_sys, dim_env)
        dm_red=np.zeros((dim_sys, dim_sys), dtype=complex)
        for i in range(dim_sys):
            for k in range(dim_sys):
                dm_red[i, k]=(psi_matrix[i].conj()*psi_matrix[k]).sum()
        #print "dm_red", dm_red
        [u, v]=np.linalg.eig(dm_red)
        #print "u left", u
        idx=u.argsort()
        u=u[idx]
        #print "u:", u.T
        #print "v left:", v
        v=v[:, idx]
        #print "v left:", v
        #print "back", np.dot(np.dot(v, np.diag(u)), np.linalg.inv(v))
        v=v.T
        #v=v.T
        v=v[-m:]
        """
        if v.sum()<0:
            v=v*(-1)
        """
        self.v=v.copy()
        #print "v left", v
        H_b1=np.kron(self.b1.H11, np.eye(self.b2.dim))
        H_sys=np.kron(self.b1.H11, np.eye(self.b2.dim))+np.kron(np.eye(self.b1.dim), self.b2.H11)
        H_t=T(self.b1, self.b2)
        H_sys=H_sys+H_t
        #print "blub", H_t
        H_sys_t=np.dot(np.dot(v.conj(), H_sys), v.T)
        new.H11=H_sys_t
        #print "new.H", new.H11
        new.szl=np.dot(np.dot(v.conj(), np.kron(np.eye(self.b1.dim), self.b2.szl)), v.T)
        new.spl=np.dot(np.dot(v.conj(), np.kron(np.eye(self.b1.dim), self.b2.spl)), v.T)
        new.sml=np.dot(np.dot(v.conj(), np.kron(np.eye(self.b1.dim), self.b2.sml)), v.T)
        #print "szl, spl, sml", new.szl, new.spl, new.sml
        new.dim=len(new.H11)
        new.sites=self.b1.sites+1
        
   
    def basis_right(self, new):
        """
        same as basis_left, only that it adds a site to the right
        """
        dim_sys=self.b3.dim*self.b4.dim
        dim_env=self.b1.dim*self.b2.dim
        psi_matrix=self.psi.reshape(dim_env, dim_sys).T
        dm_red=np.zeros((dim_sys, dim_sys), dtype=complex)
        for i in range(dim_sys):
            for k in range(dim_sys):
                dm_red[i, k]=(psi_matrix[i].conj()*psi_matrix[k]).sum()
        #print "dm:", dm_red[0, 0]
        [u, v]=np.linalg.eig(dm_red)
        idx=u.argsort()
        u=u[idx]
        #print "u right:", u.T
        #print "v right:", v.T
        v=v[:, idx].T
        v=v[-m:]
        """
        if v.sum()<0:
            v=v*(-1)
        """
        self.v=v.copy()
        H_sys=np.kron(self.b3.H11, np.eye(self.b4.dim))+np.kron(np.eye(self.b3.dim), self.b4.H11)
        H_sys+=T(self.b3, self.b4)
        H_sys_t=np.dot(np.dot(v.conj(), H_sys), v.T)
        new.H11=H_sys_t
        new.sz0=np.dot(np.dot(v.conj(), np.kron(self.b3.sz0, np.eye(self.b4.dim))), v.T)
        new.sp0=np.dot(np.dot(v.conj(), np.kron(self.b3.sp0, np.eye(self.b4.dim))), v.T)
        new.sm0=np.dot(np.dot(v.conj(), np.kron(self.b3.sm0, np.eye(self.b4.dim))), v.T)
        new.dim=len(new.H11)
        new.sites=self.b4.sites+1
        
  
        
def new_left(S):
    """
    makes new block with length l+1
    """
    res=block(1)
    S.basis_left(res)
    return res

def new_right(S):
    res=block(1)
    S.basis_right(res)
    return res


def main():
    """
    length=raw_input("Length:")
    length=int(length)
    #length=10
    nsweeps=raw_input("Number of iterations:")
    nsweeps=int(nsweeps)
    m=raw_input("m:")
    m=int(m)
    n0=raw_input("n0:")
    n0=int(n0)
    """
    blub=[]
    bla=[]
    illi=[]
    length=30
    nsweeps=2
    m=8
    n0=1
    ener=[]
    site_l=[]
    site_r=[]
    leftblock=block(n0)
    rightblock=block(n0)
    siteblock=block(1)
    all_right=[]
    all_left=[]
    site=[]
    #nsweeps=2
    
    # Warmup sweep:
    all_right.append(rightblock)
    all_left.append(leftblock)
    for i in range(int(length/2-n0-1)): 
        blub.append(i)
        bla.append(i)
        S=system()
        illi.append(i)
        S.b1=all_left[i]
        S.b2=siteblock
        S.b3=siteblock
        S.b4=all_right[i]
        site_l.append(S.b1.sites)
        site_r.append(S.b4.sites)
        site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
        S.get_gs()
        #print "psi, energy:", S.psi, S.energy
        print "warmup", i, S.energy
        ener.append(S.energy)
        bl_left=new_left(S)
        #print "H", bl_left.H11[0, :]
        #print "H_left", bl_left.H11
        #print 'bl_left', bl_left
        bl_right=new_right(S)
        all_left.append(bl_left)
        all_right.append(bl_right)
    # goes further left from the middle:
    for i in range(int(length/2-n0-1)):
        bla.append(i)
        i_l=int(i+length/2-n0-1)
        i_r=int(length/2-n0-1-i)
        blub.append(i_l)
        illi.append(i_r)
        S=system()
        S.b1=all_left[i_l]
        S.b2=siteblock
        S.b3=siteblock
        S.b4=all_right[i_r]
        site_l.append(S.b1.sites)
        site_r.append(S.b4.sites)
        site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
        S.get_gs()
        print "1/2 sweep left", i, S.energy
        ener.append(S.energy)
        bl=new_left(S)
        all_left.append(bl)
    #first full sweep right:
    #part 1: first half of H_r are already defined
    for i in range(int(length/2-n0-1)):
        bla.append(i)
        i_r=i
        i_l=int(length-2*n0-2-i)
        blub.append(i_l)
        illi.append(i_r)
        S=system()
        S.b1=all_left[i_l]
        S.b2=siteblock
        S.b3=siteblock
        S.b4=all_right[i_r]
        site_l.append(S.b1.sites)
        site_r.append(S.b4.sites)
        site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
        S.get_gs()
        print "part 1", i, S.energy
        ener.append(S.energy)
        bl=new_right(S)
        all_right[i_r+1]=bl
    #part 2: H_r need to be defined:
    for i in range(int(length/2-n0-1)):
        bla.append(i)
        i_r=int(i+length/2-n0-1)
        i_l=int(length/2-n0-1-i)
        blub.append(i_l)
        illi.append(i_r)
        S=system()
        S.b1=all_left[i_l]
        S.b2=siteblock
        S.b3=siteblock
        S.b4=all_right[i_r]
        site_l.append(S.b1.sites)
        site_r.append(S.b4.sites)
        site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
        S.get_gs()
        print "part 2", i, S.energy
        ener.append(S.energy)
        bl=new_right(S)
        all_right.append(bl)
    #finite system sweep:
    print "length all_left:", len(all_left)
    for swp in range(nsweeps):
        for i in range(int(length-2*n0-2)):
            bla.append(i)
            i_r=int(length-i-2*n0-2)
            i_l=i
            blub.append(i_l)
            illi.append(i_r)
            #print "sweep", swp, "number", i, "ir, rl:", i_r, i_l
            S=system()
            S.b1=all_left[i_l]
            S.b2=siteblock
            S.b3=siteblock
            S.b4=all_right[i_r]
            site_l.append(S.b1.sites)
            site_r.append(S.b4.sites)
            site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
            S.get_gs()
            print "sweep", swp, "left", i, S.energy
            ener.append(S.energy)
            bl=new_left(S)
            all_left[i_l+1]=bl
        for i in range(length-2*n0-2):
            bla.append(i)
            i_r=i
            i_l=length-2*n0-2-i
            blub.append(i_l)
            illi.append(i_r)
            #print "sweep", swp, "number", i, "ir, rl:", i_r, i_l
            S=system()
            S.b1=all_left[i_l]
            S.b2=siteblock
            S.b3=siteblock
            S.b4=all_right[i_r]
            site_l.append(S.b1.sites)
            site_r.append(S.b4.sites)
            site.append(S.b1.sites+S.b2.sites+S.b3.sites+S.b4.sites)
            S.get_gs()
            print "sweep", swp, "right", i, S.energy
            ener.append(S.energy)
            bl=new_right(S)
            all_right[i_r+1]=bl
    print "Energy per site:", float(ener[-1]/length)
    plt.plot(blub, ener)
    #plt.plot(site_l)
    #plt.plot(site_r)
    #ran=range(len(blub))
    #plt.plot(ran, blub)
    #plt.plot(ran, illi)
    #ges=[]
    #for i in ran:
    #    ges.append(blub[i]+illi[i])
    #plt.plot(site)
    plt.show()
    return 0

main()
  
    
